# -*- coding: UTF-8 -*-

from pymongo import MongoClient
from WindPy import *

import time
import datetime
import math



class DailyClose:
    def __init__(self, startDateStr, endDateStr = time.strftime('%Y-%m-%d'), fillType='F', DBIP='localhost', DBPort=27017, DBName = 'stocks', collectionPreName = 'daily_close', txtFileName = 'close.txt'):
        self.startDateStr = startDateStr
        self.endDateStr = endDateStr
        self.fillType = fillType
        self.DBIP = DBIP
        self.DBPort = DBPort
        self.DBName = DBName
        self.collectionPreName = collectionPreName
        self.txtFileName = txtFileName



    def writeFromWindToMongoDB(self, stockExtractGrpNum = 5):

        ###############################################################################################################
        # format of tje default 'startDateStr' is '2012-01-09'
        # format of tje default 'endDateStr' is '2017-05-08'
        # fillType = 'F': 前复权 ，fillType = 'B': 后复权
        # In order to make data extraction from WindDB faster, it's better to extract all stocks in one Wind retrieval,
        # however too large wind retrieval result causes error, therefore we split stock in to several group,
        # the variable stockExtractGrpNum indicates the group number, the default value is 5 in our case
        ###############################################################################################################
        startDateStr = self.startDateStr
        endDateStr = self.endDateStr
        fillType = self.fillType
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        collectionPreName = self.collectionPreName

        resStockID = w.wset('sectorconstituent', 'date=' + endDateStr + ';sectorid=a001010100000000;field=wind_code')

        stockIDdata = resStockID.Data[0]

        totalStockIdNum = len(stockIDdata)

        ###############################################################################################################
        #this shows stock number to retrieve in one group
        ###############################################################################################################

        splitStocksNum = math.floor(totalStockIdNum / stockExtractGrpNum)

        client = MongoClient(DBIP, DBPort)

        db = client[DBName]

        collectionName = collectionPreName + '_' + startDateStr + '-' + endDateStr + '_' + fillType

        collection_dailyclose = db[collectionName]


        for i in range(0, stockExtractGrpNum):
            stockIDOneLoopList = []
            for j in range(0, splitStocksNum):
                stockIDOneLoopList.append(i*splitStocksNum + j)
            oneStockSearchStr = ''
            for aStockID in stockIDOneLoopList:
                oneStockSearchStr = oneStockSearchStr + str(stockIDdata[aStockID]) + ', '
            resDailyClose = w.wsd(oneStockSearchStr, 'close', startDateStr, endDateStr, 'Fill=Previous;PriceAdj=' + fillType)

            batchStockNum = len(resDailyClose.Data)
            closeDataTime = resDailyClose.Times

            for k in range(0, batchStockNum):
                aData = []
                aDateDict = {}
                aDateDict['StockID'] = stockIDdata[i * splitStocksNum + k]
                totalTransDaysNum = len(closeDataTime)
                for l in range(0, totalTransDaysNum):
                    aDateDict[closeDataTime[l].strftime('%Y%m%d')] = resDailyClose.Data[k][l]
                aData.append(aDateDict)
                collection_dailyclose.insert(aData)
            print('Finished writing stocks of:' + oneStockSearchStr)

        ###############################################################################################################
        # After the looped retrieval, there could be still some stocks not retrieved, i.e., looped retrieved stocks may
        # not be the full set of available stocks, therefore we should locate the unretrieved stocks
        # by using mod function, we can check if remaining stocks is the full set of the available stock set
        # if not, we locate the remains
        ###############################################################################################################
        if (totalStockIdNum % stockExtractGrpNum) != 0:
            stockIDRemainsList = []
            for m in range((stockExtractGrpNum-1)*splitStocksNum + splitStocksNum, totalStockIdNum):
                stockIDRemainsList.append(m)
            oneStockSearchStr = ''

            for aStockID in stockIDRemainsList:
                oneStockSearchStr = oneStockSearchStr + str(stockIDdata[aStockID]) + ', '
            resDailyClose = w.wsd(oneStockSearchStr, 'close', startDateStr, endDateStr, 'Fill=Previous;PriceAdj=' + fillType)
            batchStockNum = len(resDailyClose.Data)
            closeDataTime = resDailyClose.Times

            for k in range(0, batchStockNum):
                aData = []
                aDateDict = {}
                aDateDict['StockID'] = stockIDdata[(stockExtractGrpNum-1)*splitStocksNum + splitStocksNum-1 + k]
                for l in range(0, totalTransDaysNum):
                    aDateDict[closeDataTime[l].strftime('%Y%m%d')] = resDailyClose.Data[k][l]
                aData.append(aDateDict)
                collection_dailyclose.insert(aData)
            print('Finished writing remaining stocks of:' + oneStockSearchStr)

        ###############################################################################################################
        #now we add stock index data into MongoDB
        ###############################################################################################################

        stockIndexIdList = ['000001.SH', '399001.SZ', '399006.SZ', '000300.SH', '000016.SH', '000905.SH', '399005.SZ', '881001.WI']
        stockIndexSearchStr = ''
        stockIndexNum = len(stockIndexIdList)
        for n in range(0, stockIndexNum):
            stockIndexSearchStr = stockIndexSearchStr + stockIndexIdList[n] + ', '
        resDailyClose = w.wsd(stockIndexSearchStr, 'close', startDateStr, endDateStr, 'Fill=Previous;PriceAdj=' + fillType)
        for n in range(0, stockIndexNum):
            aData = []
            aDateDict = {}
            aDateDict['StockID'] = stockIndexIdList[n]
            for l in range(0, totalTransDaysNum):
                aDateDict[closeDataTime[l].strftime('%Y%m%d')] = resDailyClose.Data[n][l]
            aData.append(aDateDict)
            collection_dailyclose.insert(aData)

        print('Finished writing stock indices of:' + stockIndexSearchStr)




    def writeFromMongoDBtoTxt(self):
        ###############################################################################################################
        # format of tje default 'startDateStr' is '2012-01-09'
        # format of tje default 'endDateStr' is '2017-05-08'
        # fillType = 'F': 前复权 ，fillType = 'B': 后复权
        ###############################################################################################################

        startDateStr = self.startDateStr
        endDateStr = self.endDateStr
        fillType = self.fillType
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        collectionPreName = self.collectionPreName
        txtFileName = self.txtFileName

        client = MongoClient(DBIP, DBPort)

        db = client[DBName]

        collectionName = collectionPreName + '_' + startDateStr + '-' + endDateStr + '_' + fillType

        collection_dailyclose = db[collectionName]

        dateList = []
        totalStockCloseList = []

        dateListRes = collection_dailyclose.find({'StockID': '000001.SH'})
        dateListDict = dateListRes[0]

        for (k, v) in dateListDict.items():
            dateStr = ''
            try:
                dateStr = datetime.datetime.strptime(k, '%Y-%m-%d').strftime('%Y%m%d')
            except:
                dateStr = k
            dateList.append(dateStr)
        dateList = dateList[2:]
        dateList.insert(0, 'close')
        totalStockCloseList.append(dateList)

        for stockIdRes in collection_dailyclose.find({}, {'StockID': 1}):
            stockIdStr = stockIdRes['StockID']
            for stockDailyCloseResDict in collection_dailyclose.find({'StockID': stockIdStr}):
                oneStockCloseList = []
                for i in stockDailyCloseResDict:
                    oneStockCloseList.append(stockDailyCloseResDict[i])
                oneStockCloseList = oneStockCloseList[1:]
                totalStockCloseList.append(oneStockCloseList)
        print('Finished retrieving close data from MongoDB...')
        f = open(txtFileName, 'w')

        txtFileRowNum = len(oneStockCloseList)
        txtFileColNum = len(totalStockCloseList)

        f.write(str(txtFileRowNum - 1) + '\t' + str(txtFileColNum - 1) + '\t')

        for firstRowLen in range(txtFileColNum - 1):
            f.write('\t')
        f.write('\n')
        print('Finished writing text file header...')

        for rowIter in range(txtFileRowNum):
            for colIter in range(txtFileColNum):
                try:
                    cellData = totalStockCloseList[colIter][rowIter]
                    if math.isnan(float(cellData)):
                        cellData = 0
                    f.write(str(cellData) + '\t')
                except ValueError:
                    cellData = str(totalStockCloseList[colIter][rowIter])
                    f.write(cellData + '\t')

            f.write('\n')
            print('Wrote row number # ' + str(rowIter) + ' to the text file....')

        f.close()
        print('Finished writing the whole text file!')

obj = DailyClose(startDateStr='2012-01-01')

w.start()
obj.writeFromWindToMongoDB()
obj.writeFromMongoDBtoTxt()
w.stop()